// VirtualMonitorUser.c
// 简单的用户态测试程序，用于与驱动通信并触发添加虚拟显示器命令
// 编译：使用 Visual Studio 的开发人员命令提示符（x64），例如：
// cl /nologo /W3 /EHsc VirtualMonitorUser.c

#include <windows.h>
#include <stdio.h>
#include <winioctl.h>

#define DEVICE_NAME L"\\\\.\\WddmFilterCtrlDevice"

// Keep compatibility with driver definitions in ioctl.h
#define IOCTL_WDDM_FILTER_CMD CTL_CODE(FILE_DEVICE_VIDEO, 0x810, METHOD_BUFFERED, FILE_ANY_ACCESS)
#define IOCTL_WDDM_FILTER_REMOVE CTL_CODE(FILE_DEVICE_VIDEO, 0x811, METHOD_BUFFERED, FILE_ANY_ACCESS)
#define IOCTL_WDDM_FILTER_ADD CTL_CODE(FILE_DEVICE_VIDEO, 0x812, METHOD_BUFFERED, FILE_ANY_ACCESS)
#define IOCTL_WDDM_FILTER_UNLOAD CTL_CODE(FILE_DEVICE_VIDEO, 0x813, METHOD_BUFFERED, FILE_ANY_ACCESS)

// Structure must match kernel's WDDM_FILTER_ADD_INPUT
#define WDDM_FILTER_MAX_NAME 64
typedef struct _WDDM_FILTER_ADD_INPUT {
    ULONG Width;
    ULONG Height;
    CHAR  Name[WDDM_FILTER_MAX_NAME];
} WDDM_FILTER_ADD_INPUT;

static void usage(const wchar_t *prog)
{
    wprintf(L"Usage:\n");
    wprintf(L"  %s --add [width height name]   (add virtual monitor, defaults used if omitted)\n", prog);
    wprintf(L"  %s --remove <id>             (remove virtual monitor by id, id=0 removes last)\n", prog);
    wprintf(L"  %s --unload                   (ask driver to unload asynchronously)\n", prog);
    wprintf(L"  %s --echo <message>           (echo message via driver)\n", prog);
}

int wmain(int argc, wchar_t **argv)
{
    HANDLE h = CreateFileW(DEVICE_NAME,
                            GENERIC_READ | GENERIC_WRITE,
                            0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (h == INVALID_HANDLE_VALUE) {
        wprintf(L"CreateFile failed: 0x%08X\n", GetLastError());
        return 1;
    }
    // parse command line
    DWORD bytesReturned = 0;
    ULONG newId = 0;
    BOOL ok = FALSE;
    if (argc >= 2) {
        // convert argv[1] to narrow for strcmp comparisons where needed
        if (wcscmp(argv[1], L"--unload") == 0) {
            ULONG out = 0;
            BOOL ok = DeviceIoControl(h, IOCTL_WDDM_FILTER_UNLOAD,
                                      NULL, 0,
                                      &out, (DWORD)sizeof(out),
                                      &bytesReturned, NULL);
            if (!ok) {
                wprintf(L"DeviceIoControl(IOCTL_WDDM_FILTER_UNLOAD) failed: 0x%08X\n", GetLastError());
            } else {
                wprintf(L"IOCTL_WDDM_FILTER_UNLOAD queued, result=%u\n", out);
            }
            CloseHandle(h);
            return 0;
        } else if (wcscmp(argv[1], L"--remove") == 0 && argc >= 3) {
            ULONG removeId = (ULONG)wcstoul(argv[2], NULL, 0);
            ULONG removeOut = 0;
            BOOL r = DeviceIoControl(h, IOCTL_WDDM_FILTER_REMOVE,
                &removeId, sizeof(removeId),
                &removeOut, sizeof(removeOut),
                &bytesReturned, NULL);
            if (r) {
                wprintf(L"REMOVE_VMON succeeded for id=%u, out=%u\n", removeId, removeOut);
            } else {
                wprintf(L"REMOVE_VMON failed: 0x%08X\n", GetLastError());
            }
            CloseHandle(h);
            return 0;
        } else if (wcscmp(argv[1], L"--add") == 0) {
            // optional parameters: width height name
            WDDM_FILTER_ADD_INPUT addIn;
            RtlSecureZeroMemory(&addIn, sizeof(addIn));
            addIn.Width = 1920;
            addIn.Height = 1080;
            strncpy_s(addIn.Name, sizeof(addIn.Name), "MyVirtualMon", _TRUNCATE);
            if (argc >= 4) {
                addIn.Width = (ULONG)wcstoul(argv[2], NULL, 0);
                addIn.Height = (ULONG)wcstoul(argv[3], NULL, 0);
            }
            if (argc >= 5) {
                // argv[4] is wide char; convert to ANSI
                size_t outLen = 0;
                wcstombs_s(&outLen, addIn.Name, sizeof(addIn.Name), argv[4], _TRUNCATE);
            }
            ok = DeviceIoControl(h, IOCTL_WDDM_FILTER_ADD,
                                      &addIn, (DWORD)sizeof(addIn),
                                      &newId, (DWORD)sizeof(newId),
                                      &bytesReturned, NULL);
            if (!ok) {
                wprintf(L"DeviceIoControl(IOCTL_WDDM_FILTER_ADD) failed: 0x%08X\n", GetLastError());
            } else {
                wprintf(L"IOCTL_WDDM_FILTER_ADD succeeded, new id=%u, bytes=%u\n", newId, bytesReturned);
            }
            // fall through to close handle and exit
        } else if (wcscmp(argv[1], L"--echo") == 0 && argc >= 3) {
            // echo message
            char mbuf[256] = {0};
            size_t outLen = 0;
            wcstombs_s(&outLen, mbuf, sizeof(mbuf), argv[2], _TRUNCATE);
            char out[512] = {0};
            ok = DeviceIoControl(h, IOCTL_WDDM_FILTER_CMD,
                                      mbuf, (DWORD)strlen(mbuf)+1,
                                      out, (DWORD)sizeof(out),
                                      &bytesReturned, NULL);
            if (ok) {
                wprintf(L"Echo succeeded, bytes=%u, echoed='%S'\n", bytesReturned, out);
            } else {
                wprintf(L"Echo failed: 0x%08X\n", GetLastError());
            }
            CloseHandle(h);
            return 0;
        } else {
            usage(argv[0]);
            CloseHandle(h);
            return 1;
        }
    } else {
        // default: add a virtual monitor with defaults
        WDDM_FILTER_ADD_INPUT addIn;
        RtlSecureZeroMemory(&addIn, sizeof(addIn));
        addIn.Width = 1920;
        addIn.Height = 1080;
        // ANSI 名称，确保 NUL 终止
        strncpy_s(addIn.Name, sizeof(addIn.Name), "MyVirtualMon", _TRUNCATE);

        ok = DeviceIoControl(h, IOCTL_WDDM_FILTER_ADD,
                                &addIn, (DWORD)sizeof(addIn),
                                &newId, (DWORD)sizeof(newId),
                                &bytesReturned, NULL);
        if (!ok) {
            wprintf(L"DeviceIoControl(IOCTL_WDDM_FILTER_ADD) failed: 0x%08X\n", GetLastError());
        } else {
            wprintf(L"IOCTL_WDDM_FILTER_ADD succeeded, new id=%u, bytes=%u\n", newId, bytesReturned);
        }
    }

    // 回显测试
    const char msg[] = "HelloKernel";
    char out[64] = {0};
    ok = DeviceIoControl(h, IOCTL_WDDM_FILTER_CMD,
                        (LPVOID)msg, (DWORD)sizeof(msg),
                        out, (DWORD)sizeof(out),
                        &bytesReturned, NULL);
    if (ok) {
        wprintf(L"Echo succeeded, bytes=%u, echoed='%S'\n", bytesReturned, out);
    } else {
        wprintf(L"Echo failed: 0x%08X\n", GetLastError());
    }

    // 演示如何删除刚才添加的虚拟显示器
    if (newId != 0) {
        ULONG removeId = newId; // 删除特定 id；若设为 0 则删除最后一个
        ULONG removeOut = 0;
        BOOL r = DeviceIoControl(h, IOCTL_WDDM_FILTER_REMOVE,
            &removeId, sizeof(removeId),
            &removeOut, sizeof(removeOut),
            &bytesReturned, NULL);
        if (r) {
            wprintf(L"REMOVE_VMON succeeded for id=%u, out=%u\n", removeId, removeOut);
        } else {
            wprintf(L"REMOVE_VMON failed: 0x%08X\n", GetLastError());
        }
    }

    CloseHandle(h);
    return 0;
}

